package grpc_proxy_middleware

import (
	"fmt"
	"gitee.com/zhoulvvv/my_go_gateway/dao"
	"gitee.com/zhoulvvv/my_go_gateway/public"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/peer"
	"log"
	"strings"
)

// GRPC 服务端和客户端限流中间件
func GRPCFlowLimitModeMiddleware(service *dao.ServiceDetail) func(srv interface{}, ss grpc.ServerStream,
	info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		serverQps := float64(service.AccessControl.ServiceFlowLimit)
		// 如果服务端限流QPS是0就表示无限制
		if serverQps > 0 {
			// 服务端限流
			serverLimiter, err := public.FlowLimiterHandler.GetLimiter(
				public.FlowServicePrefix+service.Info.ServiceName,
				serverQps)
			if err != nil {
				return errors.New("【GRPC】获取服务端限流器失败：" + err.Error())
			}
			if !serverLimiter.Allow() {
				return errors.New(fmt.Sprintf("【GRPC】服务端限流中，每秒最大QPS：%v", serverQps))
			}
		}
		clientQps := float64(service.AccessControl.ClientIPFlowLimit)
		// 如果客户端限流QPS是0就表示无限制
		if clientQps > 0 {
			peerCtx, ok := peer.FromContext(ss.Context())
			if !ok {
				log.Printf("【GRPC 服务端和客户端限流中间件中获取客户端IP失败】peer.FromContext(ss.Context()) failed with error")
			}
			splits := strings.Split(peerCtx.Addr.String(), ":")
			clientIP := ""
			if len(splits) == 2 {
				clientIP = splits[0]
			}
			// 客户端限流，客户端限流是基于IP进行限流
			clientLimiter, err := public.FlowLimiterHandler.GetLimiter(
				public.FlowServicePrefix+service.Info.ServiceName+"_"+clientIP,
				clientQps)
			if err != nil {
				return errors.New("【GRPC】获取客户端限流器失败：" + err.Error())
			}
			if !clientLimiter.Allow() {
				return errors.New(fmt.Sprintf("【GRPC】客户端限流中，每秒最大QPS：%v，IP：%s", clientQps, clientIP))
			}
		}
		if err := handler(srv, ss); err != nil {
			log.Printf("【GRPC】RPC failed with error %v\n", err)
		}
		return nil
	}
}
